#include "udp.hpp"

UDP::UDP(std::string target_addr, int target_port, size_t buffer_size)
    : target_port_(target_port), buffer_size_(buffer_size), in_group_(false) {
  // Client information
  if (!target_addr.empty())
    target_addr_ = inet_addr(target_addr.c_str());
  else
    target_addr_ = INADDR_ANY;
  cli_.sin_family = AF_INET;
  cli_.sin_addr.s_addr = target_addr_;
}

UDP::~UDP() {
  if (in_group_)
    setsockopt(sock_, IPPROTO_IP, IP_DROP_MEMBERSHIP, &multi_, sizeof(multi_));
  close(sock_);
}

int UDP::bindServer(std::string &addr, int port) {
  // Creating socket file descriptor
  if ((sock_ = socket(AF_INET, SOCK_DGRAM, 0)) < 0) {
    std::cerr << "Socket creation failed" << std::endl;
    exit(EXIT_FAILURE);
  }
  // Server information
  serv_.sin_family = AF_INET;  // IPv4
  if (addr.empty())
    serv_.sin_addr.s_addr = INADDR_ANY;
  else
    serv_.sin_addr.s_addr = inet_addr(addr.c_str());
  // If server port not set, use the same port as target's
  if (port < 0)
    serv_.sin_port = htons(target_port_);
  else
    serv_.sin_port = htons(port);
  // Bind the socket with the server address
  int ret = bind(sock_, (const struct sockaddr *)&serv_, sizeof(serv_));
  if (ret < 0) {
    std::cerr << "Bind failed" << std::endl;
    return ret;
  }
  return 0;
}

int UDP::reBindServer(std::string &addr, int port) {
  close(sock_);
  // Creating socket file descriptor
  if ((sock_ = socket(AF_INET, SOCK_DGRAM, 0)) < 0) {
    std::cerr << "Socket creation failed" << std::endl;
    exit(EXIT_FAILURE);
  }
  // Server information
  serv_.sin_family = AF_INET;  // IPv4
  if (addr.empty())
    serv_.sin_addr.s_addr = INADDR_ANY;
  else
    serv_.sin_addr.s_addr = inet_addr(addr.c_str());
  // If server port not set, use the same port as target's
  if (port < 0)
    serv_.sin_port = htons(target_port_);
  else
    serv_.sin_port = htons(port);
  // Bind the socket with the server address
  int ret = bind(sock_, (const struct sockaddr *)&serv_, sizeof(serv_));
  if (ret < 0) {
    std::cerr << "Bind failed" << std::endl;
    return ret;
  }
  return 0;
}

void UDP::singleCast(std::vector<uint8_t> &buffer) {
  cli_.sin_addr.s_addr = target_addr_;
  cli_.sin_port = htons(target_port_);
  sendto(sock_, (const uint8_t *)buffer.data(), buffer.size(), MSG_CONFIRM,
         (const sockaddr *)&cli_, sizeof(cli_));
}

void UDP::recvPacket(std::vector<uint8_t> &buffer) {
  int len, len_addr = sizeof(cli_);
  buffer.resize(buffer_size_);
  len = recvfrom(sock_, (char *)buffer.data(), buffer_size_, MSG_WAITALL,
                 (struct sockaddr *)&cli_, ((socklen_t *)(&len_addr)));
  buffer.resize(len);
}

void UDP::changeTargetPort(int port) { target_port_ = port; }
void UDP::changeTargetAddr(std::string &addr) {
  target_addr_ = inet_addr(addr.c_str());
}

int UDP::joinGroup(std::string &addr) {
  int ret;
  if (serv_.sin_addr.s_addr != INADDR_ANY) {
    std::cerr << "Server address need to be default value." << std::endl;
    std::cerr << "Rebind server with an empty address." << std::endl;
    return -1;
  }
  int loop = 1;
  ret = setsockopt(sock_, IPPROTO_IP, IP_MULTICAST_LOOP, &loop, sizeof(loop));
  if (ret < 0) {
    std::cerr << "Set loop bacp error" << std::endl;
    return -ret;
  }
  // Set broadcast address
  multi_.imr_multiaddr.s_addr = inet_addr(addr.c_str());
  multi_.imr_interface.s_addr = htonl(INADDR_ANY);
  ret =
      setsockopt(sock_, IPPROTO_IP, IP_ADD_MEMBERSHIP, &multi_, sizeof(multi_));
  if (ret < 0) {
    std::cerr << "Join group error" << std::endl;
    return ret;
  }
  in_group_ = true;
}

void UDP::recvFromGroup(std::vector<uint8_t> &buffer) {
  int len, len_addr = sizeof(serv_);
  buffer.resize(buffer_size_);
  len = recvfrom(sock_, (char *)buffer.data(), buffer_size_, 0,
                 (struct sockaddr *)&serv_, ((socklen_t *)(&len_addr)));
  if (len < 0) {
    buffer.clear();
    std::cerr << "Recv from group error" << std::endl;
    return;
  }
  buffer.resize(len);
}

void UDP::groupCast() {}

void UDP::quitGroup() {
  int err = setsockopt(sock_, IPPROTO_IP, IP_DROP_MEMBERSHIP, &multi_,
                       sizeof(multi_));
  in_group_ = false;
}
